import os
import subprocess
from flask import current_app as app


def install_package(package_name):
    subprocess.run(['pip', 'install', package_name], check=True)


def initialize_models():
    try:
        import sklearn
    except ImportError:
        app.logger.error("sklearn库不可用，尝试自动安装")
        install_package('scikit-learn')

    try:
        import jieba
    except ImportError:
        app.logger.error("jieba库不可用，尝试自动安装")
        install_package('jieba')

    model_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'services', 'text')
    os.makedirs(model_dir, exist_ok=True)

    from app.services.ocr_engine.ai_model import initialize
    initialize()


def register_retrain_route(app):
    from flask import jsonify

    @app.route('/retrain_model', methods=['POST', 'GET'])
    def retrain_model():
        from app.services.ocr_engine.ai_model import train_model
        try:
            app.logger.info("开始重新训练AI模型...")
            success = train_model()
            if success:
                app.logger.info("模型重新训练完成")
                return jsonify({"code": 1, "msg": "模型重新训练成功"})
            else:
                app.logger.error("模型重新训练失败")
                return jsonify({"code": 0, "msg": "模型重新训练失败"})
        except Exception as e:
            app.logger.error(f"重新训练失败: {str(e)}")
            return jsonify({"code": 0, "msg": f"重新训练失败: {str(e)}"})
